import base64
import datetime
import json

from skyeye_token import gzip_str, TSSHttpClient


def get_alarm_log(api_url, start_datetime, end_datetime):
    """
            :param username: 用户名
            :param login_key: 分析平台-账号管理-访问配置-单点登录密钥配置
            :param host: 分析平台IP
            :param uri: 分析平台接口
            :param params: 请求参数
            :return: 
    """
    try:
        c = TSSHttpClient(username="xxsdongj", login_key="VKGD6u3QnrgxXHsw", host="10.206.0.73")
        method = "get"
        data = {}
        gzipped_bytes = gzip_str("10.207.63.241")
        alarm_sip = base64.b64encode(gzipped_bytes)
        params = {
            "start_time": start_datetime,
            "end_time": end_datetime,
            "data_source": 1,
            "offset": 1,
            "limit": 10,
            "csrf_token": c.csrf_token
        }
        res = c.request(method=method, uri=api_url, params=params, data=data)
        result_dict = json.loads(res.text.encode('utf-8'))
        all_count = result_dict['data']['total']
        if all_count > 0:
            page_count, b = divmod(all_count, 10)
        if b > 0:
            page_count += 1
        field_list = ['access_time', 'attack_sip', 'sport', 'dip', 'dport','alarm_sip', 'hazard_level', 'host_state',
                      'ioc', 'proto', 'super_type', 'threat_name', 'type', 'repeat_count', 'location',
                      'name_type_chain', 'attack_dimension']
        for offset in range(1, page_count + 1):
            params = {
                "start_time": start_datetime,
                "end_time": end_datetime,
                "data_source": 1,
                "is_accurate": 0,
                "is_alarm_list": 1,
                "total_flag": 1,
                "offset": offset,
                "limit": 10,
                "csrf_token": c.csrf_token
            }
            res = c.request(method=method, uri=api_url, params=params, data=data)
            result_dict = json.loads(res.text.encode('utf-8'))
            for data_dict in result_dict['data']['items']:
                temp_dict = {}
                for k, v in data_dict.items():
                    if k in field_list:
                        if (isinstance(v, list)):
                            temp_dict[k] = ','.join(v)
                        else:
                            temp_dict[k] = v
                temp_dict['access_time'] = datetime.datetime.fromtimestamp(temp_dict['access_time'] / 1000)
                yield temp_dict
    except Exception as e:
        print(e)


def get_alarm_count(api_url, start_datetime, end_datetime):
    """
            :param username: 用户名
            :param login_key: 分析平台-账号管理-访问配置-单点登录密钥配置
            :param host: 分析平台IP
            :param uri: 分析平台接口
            :param params: 请求参数
            :return:
    """
    try:
        c = TSSHttpClient(username="xxsdongj", login_key="VKGD6u3QnrgxXHsw", host="10.206.0.73")
        method = "get"
        data = {}
        params = {
            "start_time": start_datetime,
            "end_time": end_datetime,
            "data_source": 1,
            "is_alarm_list": 1,
            "csrf_token": c.csrf_token
        }
        res = c.request(method=method, uri=api_url, params=params, data=data)

        result_dict = json.loads(res.text.encode('utf-8'))
        data_dict = {key: value for key, value in result_dict['data']['items'].items()}
        return data_dict
    except Exception as e:
        print(e)


import os
import sys
import django

base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(base_dir)
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'mobiledj_Security.settings')
django.setup()

from web import models


def set_datetime():
    today = datetime.datetime.now()
    current_time = today.time()
    specific_time = datetime.time(hour=15, minute=10, second=0, microsecond=0)
    if current_time > specific_time:
        create_datetime = today + datetime.timedelta(days=1)
        start_datetime = datetime.datetime(year=today.year, month=today.month, day=today.day, hour=15, minute=0,
                                           second=0)
        end_datetime = datetime.datetime(year=create_datetime.year, month=create_datetime.month,
                                         day=create_datetime.day, hour=15, minute=0, second=0)
    else:
        create_datetime = today
        last_datetime = today - datetime.timedelta(days=1)
        start_datetime = datetime.datetime(year=last_datetime.year, month=last_datetime.month, day=last_datetime.day,
                                           hour=15, minute=0, second=0)
        end_datetime = datetime.datetime(year=create_datetime.year, month=create_datetime.month,
                                         day=create_datetime.day, hour=15, minute=0, second=0)

    start_datetime_timestamp = int(start_datetime.timestamp() * 1000)
    end_datetime_timestamp = int(end_datetime.timestamp() * 1000)
    return start_datetime_timestamp, end_datetime_timestamp, create_datetime


if __name__ == '__main__':
    super_type_mapping = {
        'APT事件': 1,
        '攻击利用': 2,
        '恶意软件': 3,
        '拒绝服务': 4,
        '侦察': 5,
        '可疑活动': 6,
        '策略': 7,
        '系统': 8,
        '应用程序': 9,
        '认证': 10,
        '网络访问': 11,
    }
    hazard_level_mapping = {
        '危急': 0,
        '高危': 1,
        '中危': 2,
        '低危': 3,
    }
    host_state_mapping = {
        '企图': 1,
        '攻击成功': 2,
        '失陷': 3,
        '失败': 4,
    }
    attack_dimension_mapping = {
        '横向攻击': 1,
        '外联攻击': 2,
        '外部攻击': 3,
    }
    start_datetime, end_datetime, create_datetime = set_datetime()
    api_alarm_count = '/skyeye/v1/alarm/alarm/count'
    api_alarm_log = '/skyeye/v1/alarm/alarm/list'

    alarm_count_dict = get_alarm_count(api_alarm_count, start_datetime, end_datetime)
    alarm_count_dict.pop("unread")
    alarm_account_obj = models.SkyEye.objects.filter(create_datetime__year=create_datetime.year,
                                                     create_datetime__month=create_datetime.month,
                                                     create_datetime__day=create_datetime.day)
    if alarm_account_obj:
        alarm_count_dict['update_datetime'] = datetime.datetime.now()
        alarm_account_obj.update(**alarm_count_dict)
    else:
        alarm_count_dict['create_datetime'] = create_datetime
        models.SkyEye.objects.create(**alarm_count_dict)

    data_list = get_alarm_log(api_alarm_log, start_datetime, end_datetime)
    for row in data_list:
        obj = models.SkyEyeLog.objects.filter(
            attack_sip=row['attack_sip'],
            sport=row['sport'],
            dip=row['dip'],
            dport=row['dport'],
            name_type_chain=row['name_type_chain'],
            create_datetime__year=create_datetime.year,
            create_datetime__month=create_datetime.month,
            create_datetime__day=create_datetime.day
        )

        if not obj:
            row['create_datetime'] = create_datetime
            row['hazard_level'] = hazard_level_mapping[row['hazard_level']]
            row['super_type'] = super_type_mapping[row['super_type']]
            row['host_state'] = host_state_mapping[row['host_state']]
            row['attack_dimension'] = attack_dimension_mapping[row['attack_dimension']]
            models.SkyEyeLog.objects.create(**row)
        else:
            row['hazard_level'] = hazard_level_mapping[row['hazard_level']]
            row['super_type'] = super_type_mapping[row['super_type']]
            row['host_state'] = host_state_mapping[row['host_state']]
            row['attack_dimension'] = attack_dimension_mapping[row['attack_dimension']]
            obj.update(**row)
